import ctypes
import sys
import asyncio


ll = ctypes.cdll.LoadLibrary
lib = ll("../../InferLLM/build/libchatglm.so")


prompt_data = '四大名著有哪些'

def init(argc, argv):
    lib.init(argc, argv)

def get_remain_token():
    lib.get_remain_token.restype = ctypes.c_int
    return lib.get_remain_token()

def generate_paragraph(prompt_data):
    lib.decode.restype = ctypes.c_char_p
    output = lib.decode(prompt_data.encode('utf-8')).decode('utf-8')

    
    while(get_remain_token() > 0):
        if(lib.get_token() == lib.get_etoken()):
            print("")
            break

        lib.decode_iter.restype = ctypes.c_char_p
        o = lib.decode_iter().decode('utf-8')
        yield o
        output += o



# 准备参数
argc = len(['', '-m', '/home/khaos/workspace/chatglm-q4.bin', '-t', '8'])
argv = (ctypes.c_char_p * argc)(*map(lambda s: s.encode('utf-8'), ['', '-m', '/home/khaos/workspace/chatglm-q4.bin', '-t', '8']))


init(argc, argv)
# generate_paragraph(prompt_data)
for word in generate_paragraph(prompt_data):
    print(word)
